﻿using System;
using System.Collections.Generic;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.PathTracer;

namespace RayDen.RayEngine.Engines.MCMC
{
    public class MarkovChainSampler : ISampler
    {
        public static int MagicValue1 = 4;
        public static int MagicValue2 = 16;
        protected int MaxBSamples = 100000;

        public float LargeStepProb = 0.05f;
        public int MaxMutations = 10000;
        public int Width { get; protected set; }
        public int Height { get; protected set; }

        internal MCMCSamplerState SamplerState;

        public double[] B = new double[3];
        public double Brightness;
        private readonly object b_lock = new object();


        internal int screenWidth, screenHeight, screenStartLine = 0;
        internal int currentSampleScreenX, currentSampleScreenY, currentSubSampleIndex;
        public int pass;
        public bool lowLatency;
        public int valuesPerSample;

        public int sceneMaxPathDepth;


        private Sample TentativeSample, ChainStart;
        private float TentativeWeight;
        private int samplesInPass;
        private FastRandom rnd;
        private Stack<Sample> sampleCache;

        public MarkovChainSampler()
        {
            this.sampleCache = new Stack<Sample>(100);
            rnd = new FastRandom();
            TotalSamples = 0L;
            samplesInPass = 0;
        }
        public void Reinitialise(int seed)
        {
            this.rnd.Reinitialise(seed);
        }

        public float GetLDValue(Sample sample)
        {
            return rnd.NextFloat();

        }

        public float NextFloat()
        {
            return rnd.NextFloat();
        }
        
        

        public float PrimarySample(Sample u, int i, bool large_step, int time, int large_step_time)
        {
            if (u.modifyData[i] < time)
            {
                if (large_step)
                { // large step
                    u.Push(i); // save state
                    u.modifyData[i]  = time;
                    u.sampleData[i] = rnd.NextFloat();
                }
                else
                { // small step
                    if (u.modifyData[i] < large_step_time)
                    {
                        u.modifyData[i]  = large_step_time;
                        u.sampleData[i] = rnd.NextFloat();
                    }
                    // lazy evaluation of mutations
                    while (u.modifyData[i] < time - 1)
                    {
                        u.sampleData[i] = MC.Mutate(rnd, u.sampleData[i]);
                        u.modifyData[i]++;
                    }
                    u.Push(i); // save state
                    u.sampleData[i] = MC.Mutate(rnd, u.sampleData[i]);
                    u.modifyData[i]++;
                }
            }
            return u.sampleData[i];
        }



        protected int bSamples = 0;

        public void Add_B(ref RgbSpectrum rr)
        {
            if (ImportanceReady) 
                return;
            lock (b_lock)
            {
                if (rr.y() <= 0f)
                    return;
                bSamples++;
                
                var r = rr / rr.y();

                B[0] += r.c1;
                B[1] += r.c2;
                B[2] += r.c3;
                Brightness += rr.y();
            }
        }

        public bool ImportanceReady
        {
            get
            {
                /*
                if (bSamples == MaxBSamples)
                {
                    Brightness /= (float)MaxBSamples;
                    bSamples++;
                }    */

                return bSamples >= MaxBSamples;
            }
        }

        public RgbSpectrum Importance
        {
            get
            {
                return new RgbSpectrum((float)(B[0] / MaxMutations), (float)(B[1] / MaxMutations), (float)(B[2] / MaxMutations));
            }
        }

        public float B_I
        {
            get { return ((float)(this.Brightness /MaxMutations)); }
        }

        public Sample SmallStep(Sample sample) 
        {

            TotalSamples++;

            var s = new Sample(sample);
            s.imageX = MC.Mutate(s.imageX, rnd.NextFloat(), rnd.NextFloat());
            s.imageY = MC.Mutate(s.imageY, rnd.NextFloat(), rnd.NextFloat());

            for (int i = 0; i < s.sampleData.Length; i++)
            {
                s.sampleData[i] = Mutate(s.sampleData[i], rnd.NextFloat(), rnd.NextFloat());
            }
            return s;
        }
        protected void SmallStep()
        {

            TotalSamples++;

            var s = TentativeSample;
            s.imageX = MC.Mutate(s.imageX, rnd.NextFloat(), rnd.NextFloat());
            s.imageY = MC.Mutate(s.imageY, rnd.NextFloat(), rnd.NextFloat());

            for (int i = 0; i < s.sampleData.Length; i++)
            {
                s.sampleData[i] = Mutate(s.sampleData[i], rnd.NextFloat(), rnd.NextFloat());
            }
        }
        public Sample GetNextSample(Sample oldSample, float importance,float oldI,float b, ref RgbSpectrum contrib, out RgbSpectrum newsample_contrib, out float newW, bool largeStep)
        {
            var newSample = new Sample(oldSample);
            float a = Math.Min(1f, importance / oldI);
            newsample_contrib = contrib;
            newW = (a + (largeStep?1f:0f)) / (importance / b + LargeStepProb) / MaxMutations;

            return newSample;
        }

        
        public Sample GetSample(Sample prevSample)
        {
            TotalSamples++;
            Sample sample = new Sample(this);
            sample.sampleData = new float[SampleDataOffset.SamplesPerVertex * sceneMaxPathDepth];
            //sample.modifyData = new int[SampleDataOffset.SamplesPerVertex * sceneMaxPathDepth];
            //MC.StratifiedSample1D(sample.sampleData, SampleDataOffset.SamplesPerVertex * sceneMaxPathDepth, rnd, false);
            //sample.sampleData 
            sample.InitSample(SampleDataOffset.SamplesPerVertex * sceneMaxPathDepth);

            int scrX, scrY;

            sample.imageX = rnd.NextFloat()*screenWidth;
            sample.imageY = rnd.NextFloat()*screenHeight;
            //sample.imageY += screenStartLine;
            return sample;
            /*
            if ((pass >= MagicValue2))
            {
                // In order to improve ray coherency
                int stepX = currentSubSampleIndex % 4;
                int stepY = currentSubSampleIndex / 4;

                scrX = currentSampleScreenX;
                scrY = currentSampleScreenY;

                currentSubSampleIndex++;
                if (currentSubSampleIndex == (MagicValue1 / 2))
                {
                    currentSubSampleIndex = 0;
                    currentSampleScreenX++;
                    if (currentSampleScreenX >= screenWidth)
                    {
                        currentSampleScreenX = 0;
                        currentSampleScreenY++;

                        if (currentSampleScreenY >= screenHeight)
                        {
                            currentSampleScreenY = 0;
                            pass += 16;
                        }
                    }
                }

                float r1 = (stepX + ((float)rnd.NextDouble())) / 4f - 0.5f;
                float r2 = (stepY + ((float)rnd.NextDouble())) / 4f - 0.5f;
                sample.imageX = scrX + r1;
                sample.imageY = scrY + r2;
                //sample.Init(this,pass, valuesPerSample, false);
            }
            else if (pass >= MagicValue1)
            {
                scrX = currentSampleScreenX;
                scrY = currentSampleScreenY;

                currentSampleScreenX++;
                if (currentSampleScreenX >= screenWidth)
                {
                    currentSampleScreenX = 0;
                    currentSampleScreenY++;

                    if (currentSampleScreenY >= screenHeight)
                    {
                        currentSampleScreenY = 0;
                        pass++;
                    }
                }

                float r1 = ((float)rnd.NextDouble()) - .5f;
                float r2 = ((float)rnd.NextDouble()) - .5f;
                sample.imageX = scrX + r1;
                sample.imageY = scrY + r2;
                //sample.Init(this, pass,valuesPerSample, false);
            }
            else
            {
                // In order to update the screen faster for the first 16 passes
                for (; ; )
                {
                    int stepX = pass % 4;
                    int stepY = (pass / 4) % 4;

                    scrX = currentSampleScreenX * 4 + stepX;
                    scrY = currentSampleScreenY * 4 + stepY;

                    currentSampleScreenX++;
                    if (currentSampleScreenX * 4 >= screenWidth)
                    {
                        currentSampleScreenX = 0;
                        currentSampleScreenY++;

                        if (currentSampleScreenY * 4 >= screenHeight)
                        {
                            currentSampleScreenY = 0;
                            pass++;
                        }
                    }

                    // Check if we are inside the screen
                    if ((scrX < screenWidth) && (scrY < screenHeight))
                    {
                        // Ok, it is a valid sample
                        break;
                    }
                    else if (pass >= MagicValue1)
                    {
                        sample = this.GetSample(prevSample);
                        return sample;
                    }
                }

                float r1 = ((float)rnd.NextDouble()) - .5f;
                float r2 = ((float)rnd.NextDouble()) - .5f;

                sample.imageX = scrX + r1;
                sample.imageY = scrY + r2;
                //sample.Init(this, pass, valuesPerSample, false);
            }
            sample.imageY += screenStartLine;
            return sample;*/
        }

        public static float Mutate(float value, float u, float v)
        {
            const float s1 = 1f / 512f, s2 = 1f / 16f;
            float dv = (float)(s2 * Math.Exp(-Math.Log(s2 / s1) * u));
            if (v < 0.5f)
            {
                value += dv; if (value > 1f) value -= 1f;
            }
            else
            {
                value -= dv; if (value < 0f) value += 1f;
            }
            return value;
        }

        #region ISampler Members


        public float GetLazyValue(Sample sample)
        {
            return sample.GetLazyValue(pass);
        }

        public void Init(int width, int height)
        {
            rnd.Reinitialise((int)DateTime.Now.Ticks);
            if (screenWidth == 0 && screenHeight == 0)
            {
                Width = screenWidth = width;
                Height = screenHeight = height;
                screenWidth = width - 1;
                screenHeight = height - 1;
            }
            currentSampleScreenX = 0;
            currentSampleScreenY = screenStartLine;
            currentSubSampleIndex = 0;
            pass = 0;
            this.SamplerState = MCMCSamplerState.Initialized;
        }

        private int mutationsCount;

        protected void PushSample()
        {
            this.sampleCache.Push(new Sample(TentativeSample));
        }

        protected bool PopSample()
        {
            if (this.sampleCache.Count == 0)
                return false;
            this.TentativeSample = new Sample(this.sampleCache.Pop());
            return true;
        }

        public Sample EvalNextSample(Sample oldSample, float importance)
        {
            if (oldSample == null)
                return GetSample(null);
            switch (SamplerState)
            {
                case MCMCSamplerState.LargeStep:
                case MCMCSamplerState.Initialized:
                    if (importance > 0f)
                    {
                        this.SamplerState = MCMCSamplerState.SmallStep;
                        this.mutationsCount = 0;
                        this.TentativeSample = new Sample(oldSample);
                        this.ChainStart = new Sample(oldSample);
                        this.TentativeWeight = importance;
                        this.TotalSamples++;

                        return this.TentativeSample;
                    }
                    else
                    {
                        return this.GetSample(null);
                    }
                case MCMCSamplerState.SmallStep:
                    if (importance > 0f)
                    {
                        mutationsCount++;
                        PushSample();
                        this.SmallStep();
                    }
                    else
                    {
                        if (!PopSample())
                        {
                            this.TentativeSample = new Sample(ChainStart);
                        }                        
                    }
                    if (mutationsCount > MaxMutations || (rnd.NextFloat() < LargeStepProb))
                    {
                        this.SamplerState = MCMCSamplerState.LargeStep;
                        this.mutationsCount = 0;
                        return GetSample(null);
                    }

                    this.TotalSamples++;
                    return this.TentativeSample;
                default:
                    return GetSample(null);

            }
        }

        public int GetPass()
        {
            return pass;
        }

        public long TotalSamples { get; set; }
        public float GetLazyValue(Sample sample, int index)
        {
            return sample.sampleData[unchecked(index)];
        }


        #endregion

        internal float EvalImportance(ref Library.Core.Primitives.RgbSpectrum rgbSpectrum, float p)
        {
            return (rgbSpectrum).y();
        }
    }
}
